package com.yanyu.space.common.sleuth.web.response;

import cn.hutool.core.bean.BeanUtil;
import com.yanyu.space.common.core.bean.web.R;
import com.yanyu.space.common.sleuth.util.TracerWebUtil;
import org.springframework.cloud.sleuth.TraceContext;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

/**
 * @author yanyu
 */
@ControllerAdvice
public class TraceIdResponseBodyAdvice implements ResponseBodyAdvice<R> {
    @Override
    public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
        return R.class.isAssignableFrom(returnType.getParameterType());
    }

    @Override
    public R beforeBodyWrite(R body, MethodParameter returnType, MediaType selectedContentType, Class<? extends HttpMessageConverter<?>> selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {
        TraceContext traceContext = TracerWebUtil.getTraceContextInRequest();
        if(traceContext != null) {
            TraceR traceR = BeanUtil.copyProperties(body, TraceR.class);
            traceR.setRequestId (traceContext.traceId());
            return traceR;
        }
        return body;
    }
}
